35 research outputs found
Distributed Pruning Towards Tiny Neural Networks in Federated Learning
Neural network pruning is an essential technique for reducing the size and
complexity of deep neural networks, enabling large-scale models on devices with
limited resources. However, existing pruning approaches heavily rely on
training data for guiding the pruning strategies, making them ineffective for
federated learning over distributed and confidential datasets. Additionally,
the memory- and computation-intensive pruning process becomes infeasible for
recourse-constrained devices in federated learning. To address these
challenges, we propose FedTiny, a distributed pruning framework for federated
learning that generates specialized tiny models for memory- and
computing-constrained devices. We introduce two key modules in FedTiny to
adaptively search coarse- and finer-pruned specialized models to fit deployment
scenarios with sparse and cheap local computation. First, an adaptive batch
normalization selection module is designed to mitigate biases in pruning caused
by the heterogeneity of local data. Second, a lightweight progressive pruning
module aims to finer prune the models under strict memory and computational
budgets, allowing the pruning policy for each layer to be gradually determined
rather than evaluating the overall model structure. The experimental results
demonstrate the effectiveness of FedTiny, which outperforms state-of-the-art
approaches, particularly when compressing deep models to extremely sparse tiny
models. FedTiny achieves an accuracy improvement of 2.61% while significantly
reducing the computational cost by 95.91% and the memory footprint by 94.01%
compared to state-of-the-art methods.Comment: This paper has been accepted to ICDCS 202
DOMINO++: Domain-aware Loss Regularization for Deep Learning Generalizability
Out-of-distribution (OOD) generalization poses a serious challenge for modern
deep learning (DL). OOD data consists of test data that is significantly
different from the model's training data. DL models that perform well on
in-domain test data could struggle on OOD data. Overcoming this discrepancy is
essential to the reliable deployment of DL. Proper model calibration decreases
the number of spurious connections that are made between model features and
class outputs. Hence, calibrated DL can improve OOD generalization by only
learning features that are truly indicative of the respective classes. Previous
work proposed domain-aware model calibration (DOMINO) to improve DL
calibration, but it lacks designs for model generalizability to OOD data. In
this work, we propose DOMINO++, a dual-guidance and dynamic domain-aware loss
regularization focused on OOD generalizability. DOMINO++ integrates
expert-guided and data-guided knowledge in its regularization. Unlike DOMINO
which imposed a fixed scaling and regularization rate, DOMINO++ designs a
dynamic scaling factor and an adaptive regularization rate. Comprehensive
evaluations compare DOMINO++ with DOMINO and the baseline model for head tissue
segmentation from magnetic resonance images (MRIs) on OOD data. The OOD data
consists of synthetic noisy and rotated datasets, as well as real data using a
different MRI scanner from a separate site. DOMINO++'s superior performance
demonstrates its potential to improve the trustworthy deployment of DL on real
clinical data.Comment: 12 pages, 5 figures, 5 tables, Accepted by the International
Conference on Medical Image Computing and Computer Assisted Intervention
(MICCAI) 202